# 无人机上升目标高度最短用时控制-动态规划数值方法 （4.2节案例）
# 逆向分级求解方法
# 理论基础：贝尔曼最优化理论

# 系统描述：无人机高度控制，配合PPT图片理解
# 目标：使无人机在给定时间内从0m到达10米；h(0)=0,dh(0)/dt=0,hd=10,dhd/dt=0,
# 约束条件：加速度：[-3,2]；速度：[0,3]

import numpy as np
import math
import matplotlib.pyplot as plt
import pandas as pd

# 1.初始化无人机状态，目标状态，约束条件
# 无人机高度，速度
h_init = 0 
v_init = 0
# ※ 初速度不为0，会报错。
# 无人机终点目标：高度，速度
h_final = 10
v_final = 0
# 高度和速度和加速度约束条件
h_min = 0
h_max = 10
v_min = 0
v_max = 3
u_min = -3
# u_min = 0.1
u_max = 2
# 高度离散步长
N_h = 5
# 速度离散步长
N_v = 3
# 2.初始化存入各个数据的矩阵
# 离散向量
Hd = np.arange(h_min, h_max + (h_max - h_min) / N_h , (h_max - h_min) / N_h )
Vd = np.arange(v_min, v_max + (v_max - v_min) / N_v , (v_max - v_min) / N_v)
# Vd = np.repeat(Vd[:, np.newaxis], N_h + 1, axis=1)
# 代价函数矩阵,存入每个离散点的代价
J_costtogo = np.zeros((N_h + 1, N_v + 1))
# 存每个离散点的加速度
Input_acc = np.zeros((N_h + 1, N_v + 1))


# 3.开始逆向分级求解

# 3.1 最顶点（N_h+1）到第N_h点的代价和加速度计算
v_avg = 0.5 * (v_final + Vd)
# Calculate travel time, this is the cost
T_delta = ((h_max - h_min)/N_h)/v_avg
# 加速度计算
acc = (v_final - Vd)/T_delta
# Assign delta T to cost to go（中间代价）
J_temp = T_delta
# 查找超过约束条件的加速度的索引
Ind_lin_acc = np.where((acc < u_min) | (acc > u_max))[0]
# 根据查找到的索引，给代价函数矩阵赋值为无穷大
J_temp[Ind_lin_acc] = np.inf
# 代价将暂存值赋值给最终的代价矩阵
J_costtogo[1,:] = J_temp
# 加速度将暂存值赋值给最终的加速度矩阵
Input_acc[1,:] = acc
# # 3.2 第N_h点到第2点的代价和加速度计算
for k in range(2, N_h, 1):
    # 将vd重复N_h+1次，作为列向量和行向量
    Vd_x, Vd_y = np.meshgrid(Vd, Vd)
    # 计算平均速度
    v_avg = 0.5 * (Vd_x + Vd_y)
    # 计算行走时间
    T_delta = ((h_max - h_min)/N_h) / v_avg
    # 计算加速度
    acc = (Vd_y - Vd_x)/T_delta
    # 暂存代价
    J_temp = T_delta
    # 查找超过约束条件的加速度的索引（行列索引）
    row, col = np.where((acc < u_min) | (acc > u_max))
    # 把超出加速度约束的代价赋值无穷大
    J_temp[row,col] = np.inf
    # 把上一次的代价，以列向量的形式加到J_temp的每个列上
    x_temp, y_temp = np.meshgrid(J_costtogo[k-1,:],J_costtogo[k-1,:])
    J_temp = J_temp + y_temp
    # 找到J_temp中每一列的最小值，作为当前高度离散点的代价
    J_costtogo[k,:] = (np.min(J_temp,axis=0)).reshape(1,-1)
    # 沿着矩阵 J_temp 的 列（axis=0） 找到每一列的最小值的行索引
    row_indices = np.argmin(J_temp, axis=0)
    # 生成一个从0到J_temp.shape[1]（列数）的整数数组
    col_indices = np.arange(J_temp.shape[1])
    # 将行列索引以J_temp的形式，转化为线性索引
    linear_indices = np.ravel_multi_index((row_indices, col_indices), J_temp.shape)
    # 将线性索引转化为加速度的矩阵形式
    Input_acc [k,:] = acc.flat[linear_indices].reshape(1,-1)
   
# 3.3 第2点到第1点的代价和加速度计算
# 计算平均速度
v_avg = 0.5 * (Vd + v_init)
# 计算行走时间
T_delta = (h_max - h_min)/(N_h * v_avg)
# 计算加速度
acc = (Vd - v_init)/T_delta
# 暂存代价
J_temp = T_delta
# 查找超过约束条件的加速度的索引
Ind_lin_acc = np.where((acc < u_min) | (acc > u_max))[0]
# 根据查找到的索引，给代价函数矩阵赋值为无穷大
J_temp[Ind_lin_acc] = np.inf
# 将最后一次的代价加上上一次的代价
J_temp =J_temp + J_costtogo[N_h-1,:]
# 查找本次代价中最小的代价，作为最后一次的代价，并放入J_costtogo的首位置
J_costtogo[N_h,0] = np.min(J_temp,axis=0)
# 将最后一次最小代价位置上的加速度，赋值给Input_acc的首位置
Ind_lin_acc = np.where(J_temp==min(J_temp))[0]
Input_acc[N_h,0] = acc[Ind_lin_acc]

# 4.输出结果-画图 

h_plot_init = 0
v_plot_init = 0
acc_plot = np.zeros(len(Hd))
v_plot = np.zeros(len(Hd))
h_plot = np.zeros(len(Hd))
h_plot[0] = h_plot_init
v_plot[0] = v_plot_init
# 写入数据到表格中
# from openpyxl import Workbook
# # 创建 Excel 工作簿
# wb = Workbook()
# ws = wb.active
# # 将矩阵数据写入工作表
# for row in Input_acc:
#     ws.append(row.tolist())
# # 保存 Excel 文件
# wb.save('Input_acc.xlsx')

for k in range(0, N_h):
    diff_h = np.abs(h_plot[k] - Hd)
    # 找到最小值和索引
    min_h = np.min(diff_h)
    h_plot_index = np.argmin(diff_h)
    
    diff_v = np.abs(v_plot[k] - Vd)
    # 找到最小值和索引
    min_v = np.min(diff_v)
    v_plot_index = np.argmin(diff_v)
    
    # 将行列索引以Input_acc的形式，转化为线性索引,先列后行
    # acc_index = np.ravel_multi_index((N_h-h_plot_index, v_plot_index), Input_acc.shape, order='F')
    acc_plot[k] = Input_acc[N_h-h_plot_index, v_plot_index]

    v_plot[k+1] = math.sqrt((2 * (h_max - h_min)/N_h * acc_plot[k])+ (v_plot[k])**2)
    h_plot[k+1] = h_plot[k] + (h_max - h_min)/N_h
    

#  Plot
plt.subplot(2,1,1)
plt.plot(v_plot,h_plot,'--^')
plt.title('速度/加速度--高度')
plt.ylabel('h(m)')
plt.xlabel('v(m/s)')

plt.subplot(2,1,2)
plt.plot(acc_plot,h_plot,'^')
plt.ylabel('h(m)')
plt.xlabel('a(m/s^2)')

plt.show()